package oktx.narayana;

import com.arjuna.ats.jta.recovery.XAResourceRecoveryHelper;
import org.springframework.boot.jdbc.XADataSourceWrapper;
import org.springframework.boot.jta.narayana.DataSourceXAResourceRecoveryHelper;
import org.springframework.boot.jta.narayana.NarayanaProperties;
import org.springframework.boot.jta.narayana.NarayanaRecoveryManagerBean;
import org.springframework.util.Assert;

import javax.sql.DataSource;
import javax.sql.XADataSource;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;

/**
 * Created by huangdachao on 2018/6/5 17:01.
 */
public class NarayanaXADataSourceWrapper implements XADataSourceWrapper {
    private final NarayanaRecoveryManagerBean recoveryManager;
    private final NarayanaProperties properties;
    private final Method registerRecoveryHelperMethod;

    /**
     * Create a new {@link org.springframework.boot.jta.narayana.NarayanaXADataSourceWrapper} instance.
     * @param recoveryManager the underlying recovery manager
     * @param properties the Narayana properties
     */
    public NarayanaXADataSourceWrapper(NarayanaRecoveryManagerBean recoveryManager,
                                       NarayanaProperties properties) {
        Assert.notNull(recoveryManager, "RecoveryManager must not be null");
        Assert.notNull(properties, "Properties must not be null");
        this.recoveryManager = recoveryManager;
        this.properties = properties;
        try {
            registerRecoveryHelperMethod = NarayanaRecoveryManagerBean.class.getDeclaredMethod(
                "registerXAResourceRecoveryHelper", XAResourceRecoveryHelper.class);
            registerRecoveryHelperMethod.setAccessible(true);
        } catch (NoSuchMethodException | SecurityException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public DataSource wrapDataSource(XADataSource dataSource) {
        XAResourceRecoveryHelper recoveryHelper = getRecoveryHelper(dataSource);
        try {
            registerRecoveryHelperMethod.invoke(this.recoveryManager, recoveryHelper);
        } catch (InvocationTargetException | IllegalAccessException e) {
            throw new RuntimeException(e);
        }
        return new NarayanaDataSource(dataSource);
    }

    private XAResourceRecoveryHelper getRecoveryHelper(XADataSource dataSource) {
        if (this.properties.getRecoveryDbUser() == null
            && this.properties.getRecoveryDbPass() == null) {
            return new DataSourceXAResourceRecoveryHelper(dataSource);
        }
        return new DataSourceXAResourceRecoveryHelper(dataSource,
            this.properties.getRecoveryDbUser(), this.properties.getRecoveryDbPass());
    }
}
